import tensorflow as tf

def freeze_graph(input_checkpoint, output_graph):
    '''
    model: the defined model
    checkpoint_dir: the dir of three files
    model_name: the name of .pb
    '''
    graph = tf.Graph()
    config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
    with tf.Session(graph=graph, config=config) as sess:
        # Restore from checkpoint
        loader = tf.compat.v1.train.import_meta_graph(input_checkpoint + '.meta')
        loader.restore(sess, input_checkpoint)
        # Export checkpoint to SavedModel
        builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(output_graph)
        builder.add_meta_graph_and_variables(sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING],
                                             strip_default_attrs=True)
        builder.save()
def get_ckptname(input_checkpoint):
    from tensorflow.python import pywrap_tensorflow
    reader = pywrap_tensorflow.NewCheckpointReader(input_checkpoint)
    var_to_shape_map = reader.get_variable_to_shape_map()
    for key in var_to_shape_map:
        print("tensor_name: ", key)


if __name__ == '__main__':
    # 输入ckpt模型路径
    input_checkpoint = 'colorlization_train/model-858570'
    # 输出pb模型的路径
    out_pb_path = "frozen_model.pb"
    # 调用freeze_graph将ckpt转为pb
    # freeze_graph(input_checkpoint, out_pb_path)
    get_ckptname(input_checkpoint)
